#include <iostream>
#include <iomanip>
#include <vector>
#include <cstdlib>
#include <omp.h>
#include <complex>
#include <mambo.hpp>
#include <dns_ops.h>
using namespace std;

int main(int argc, char** argv)
{

  typedef  double real_type;
  typedef  complex<double> complex_type;
  const int n_var=3;
  int niters=2;
  int n=12288;
  int npx=512;
  int npy=512;
  int ic=0;
  while(ic<argc)
  {
    string a(argv[ic]);
    if(a.find("-s")<a.size()) //get the z
      n=atoi(argv[++ic]);
    else if(a.find("-npx")<a.size()) //x-partition
      npx=atoi(argv[++ic]);
    else if(a.find("-npy")<a.size()) //x-partition
      npy=atoi(argv[++ic]);
    else if(a.find("-i")<a.size()) 
      niters=atoi(argv[++ic]);
    else if(a.find("-h")<a.size()) 
      cerr << "Usage: -s fft-size -npx xgrid -npy ygrid -i iterations " << endl;

    ++ic;
  }

  int nx=n/npx;
  int ny=n/npy;
  int nz=n;
  int ny_half=ny/2;
  int ntot=nx*ny*nz*3;

  int np=omp_get_max_threads();
  real_type *A = new real_type[np*ntot];
  real_type *B = new real_type[np*ntot];
  int *kx_map=new int[nx];
  int *ky_map=new int[ny_half];
  int *kz_map=new int[nz];

  cout.setf(ios::fixed, ios::floatfield);
  cout.precision(5);
  cout << "BW nx_loc=" << nx << " ny_loc=" << ny << " nz_loc=" << nz
    << endl;

  cout << "PHASEA" << endl;
  //mambo::warmup();
  std::fill(kx_map,kx_map+nx,9);
  std::fill(ky_map,ky_map+ny_half,9);
  std::fill(kz_map,kz_map+nz,9);

  cout << "PHASEC" << endl;
  //mambo::begin();
  Timer wtimer;
  double t0=0,t1=0,t2=0;
//#pragma omp parallel reduction(+:t0,t1,t2)
//  {
#pragma omp parallel for reduction(+:t0,t1,t2)
  for(int ip=0;ip<np;++ip)
  {
    int first=ntot*ip;
    int last=first+ntot;
    std::fill(A+first,A+last,ip);
    std::fill(B+first,B+last,2*ip);

    double timing[]={0.0,0.0,0.0,0.0};
    Timer myclock;
    for(int i=0; i<niters; ++i)
    {
      myclock.restart();
      dns_ops::vorticity_xyz3(A+first,B+first,kx_map,nx,ky_map,ny_half,kz_map,nz);
      timing[0]+=myclock.elapsed();

      myclock.restart();
      int nxyz3=ntot/2;
      vxyz3(complex_ptr(A+first),complex_ptr(B+first),
          kx_map,&nx,ky_map,&ny_half,kz_map,&nz,&nxyz3);
      timing[1]+=myclock.elapsed();

      myclock.restart();
      //dns_ops::vorticity_xy3z(B+first,A+first,kx_map,nx,ky_map,ny_half,kz_map,nz);
      dns_ops::vorticity_3xyz(A+first,B+first,kx_map,nx,ky_map,ny_half,kz_map,nz);
      timing[2]+=myclock.elapsed();

    }

    t0=timing[0];
    t1=timing[1];
    t2=timing[2];
  }
  double tot_time=wtimer.elapsed();
  double nodefac= 1/static_cast<double>(niters*omp_get_max_threads());
  //using 4 for RK4
  cout << "BW OMP_NUM_THREADS= " << omp_get_max_threads() 
    <<  " Wall_clock= " << tot_time << " sec "
    << (t0+t1+t2)/static_cast<double>(np) 
    << endl;
  cout << "BW fft-dim num-ffts ops c-xyz3 f-xyz3 c-3xyz" << endl;
  cout << "BW " << nz << " " << nx*ny  << " " << ntot
    << " " << t0*nodefac << " " << t1*nodefac  << " " << t2*nodefac
    << endl;

  //mambo::end();
  cout << "PHASED" << endl;

  delete [] A;
  delete [] B;
  delete [] kx_map;
  delete [] ky_map;
  delete [] kz_map;

  return 0;
}

